import gzip
import os

import numpy as np


def load_mnist_images(filename):
    """
    加载MNIST的图片数据。
    """
    file_path = os.path.join('../深度学习入门/mnist_data', filename)
    with gzip.open(file_path, 'rb') as f:
        # MNIST的图片数据文件的前16个字节是文件头信息，需要跳过
        # 之后的数据才是真正的像素值
        data = np.frombuffer(f.read(), np.uint8, offset=16)
    # 将一维数组重塑为 (图片数量, 28, 28) 的三维数组
    images = data.reshape(-1, 28, 28)
    return images

def load_mnist_labels(filename):
    """
    加载MNIST的标签数据。
    """
    file_path = os.path.join('../深度学习入门/mnist_data', filename)
    with gzip.open(file_path, 'rb') as f:
        # 标签文件的前8个字节是文件头信息
        data = np.frombuffer(f.read(), np.uint8, offset=8)
    return data


def preprocess_data(images, labels):
    """
    数据预处理：
    1.  将28x28的图片“压平”成784维的向量。
    2.  将像素值从[0, 255]归一化到[0.01, 1.0]。
        (我们不缩放到[0,1]是为了避免0值输入导致权重更新出问题)
    3.  将标签进行"one-hot"编码。例如，数字5会变成 [0,0,0,0,0,1,0,0,0,0]。
    """
    # 压平图片并归一化
    images_flattened = images.reshape(images.shape[0], -1).astype(np.float32)
    images_normalized = (images_flattened / 255.0 * 0.99) + 0.01

    # One-hot编码标签
    # 创建一个全为0.01的矩阵，然后将对应标签位置的值设为0.99
    # 创建一个形状为 (样本数, 10) 的二维数组，初始值为 0.01。
    # labels.size 是标签的总个数。
    # 每一行代表一个样本的 one-hot 编码，10 是因为 MNIST 是 10 类（数字 0-9）。
    labels_one_hot = np.zeros((labels.size, 10)) + 0.01

    # 假设标签是
    # labels = [2, 5]，处理后的结果是：
    # labels_one_hot =
    # [
    #  [0.01, 0.01, 0.99, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],  # 第1个样本是类2
    #  [0.01, 0.01, 0.01, 0.01, 0.01, 0.99, 0.01, 0.01, 0.01, 0.01]   # 第2个样本是类5
    # ]
    for i, label in enumerate(labels):
        labels_one_hot[i][label] = 0.99

    return images_normalized, labels_one_hot

# 执行下载和加载
print("--- 开始加载数据 ---")
# download_mnist()
train_images = load_mnist_images('train-images-idx3-ubyte.gz')
train_labels = load_mnist_labels('train-labels-idx1-ubyte.gz')
test_images = load_mnist_images('t10k-images-idx3-ubyte.gz')
test_labels = load_mnist_labels('t10k-labels-idx1-ubyte.gz')


# 预处理数据
X_train, y_train = preprocess_data(train_images, train_labels)
X_test, y_test_one_hot = preprocess_data(test_images, test_labels)
# 测试时，我们还是用原始的数字标签进行比较
y_test = test_labels

print("数据加载与预处理完成！")
print(f"训练集大小: {X_train.shape[0]} 条")
print(f"测试集大小: {X_test.shape[0]} 条")
print("-" * 20)

